import os
import sys
import subprocess
import shutil
import urllib.request
import zipfile
import ctypes
import winreg
import requests

CUDA_VERSION = "12.9.1"
CUDA_MAJOR = "12"
CUDA_MINOR = "9"
CUDA_BASE = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA"
CUDA_PATH = os.path.join(CUDA_BASE, f"v{CUDA_MAJOR}.{CUDA_MINOR}")

# 直接用官网复制的真实链接
INSTALLER_URL = "https://developer.download.nvidia.com/compute/cuda/12.9.1/local_installers/cuda_12.9.1_576.57_windows.exe"
# 替换为你自己的 cuDNN zip 包直链
CUDNN_URL = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/windows-x86_64/cudnn_9.10.2_windows.exe"

def is_admin():
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False

def download_file(url, dest):
    print(f"[信息] 正在下载: {url}")
    headers = {"User-Agent": "Mozilla/5.0"}
    with requests.get(url, headers=headers, stream=True) as r:
        r.raise_for_status()
        total = int(r.headers.get('content-length', 0))
        with open(dest, "wb") as f:
            downloaded = 0
            for chunk in r.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    downloaded += len(chunk)
                    done = int(50 * downloaded / total)
                    print(f"\r[下载进度] [{'#' * done}{'.' * (50 - done)}] {downloaded/1024/1024:.2f}MB/{total/1024/1024:.2f}MB", end="")
    print("\n下载完成！")

def install_cuda():
    if os.path.exists(CUDA_PATH):
        print(f"[信息] 已检测到 CUDA Toolkit {CUDA_MAJOR}.{CUDA_MINOR} 已安装于: {CUDA_PATH}")
        return
    installer_path = os.path.join(os.environ['TEMP'], "cuda_12.9.1_576.57_windows.exe")
    if os.path.exists(installer_path):
        os.remove(installer_path)
    download_file(INSTALLER_URL, installer_path)
    print("[信息] 正在静默安装 CUDA Toolkit，请耐心等待...")
    ret = subprocess.call([installer_path, "-s"])
    if ret != 0:
        print("[错误] CUDA Toolkit 安装失败！")
        sys.exit(1)
    if not os.path.exists(CUDA_PATH):
        print(f"[错误] CUDA Toolkit 安装后未检测到目录: {CUDA_PATH}")
        sys.exit(1)
    print("[信息] CUDA Toolkit 安装成功！")

def install_cudnn():
    cudnn_zip = os.path.join(os.environ['TEMP'], "cudnn.zip")
    if os.path.exists(cudnn_zip):
        os.remove(cudnn_zip)
    download_file(CUDNN_URL, cudnn_zip)
    extract_dir = os.path.join(os.environ['TEMP'], "cudnn_tmp")
    if os.path.exists(extract_dir):
        shutil.rmtree(extract_dir)
    print("[信息] 正在解压 cuDNN ...")
    with zipfile.ZipFile(cudnn_zip, 'r') as zip_ref:
        zip_ref.extractall(extract_dir)
    # 复制文件
    for subdir in ['bin', 'include', os.path.join('lib', 'x64')]:
        src = os.path.join(extract_dir, 'cuda', subdir)
        dst = os.path.join(CUDA_PATH, subdir)
        if os.path.exists(src):
            if not os.path.exists(dst):
                os.makedirs(dst)
            for file in os.listdir(src):
                shutil.copy2(os.path.join(src, file), dst)
    shutil.rmtree(extract_dir)
    os.remove(cudnn_zip)
    print("[信息] cuDNN 安装完成！")

def set_env():
    print("[信息] 正在配置环境变量...")
    subprocess.call(f'setx CUDA_PATH "{CUDA_PATH}"', shell=True)
    reg_path = r"SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment"
    with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, reg_path, 0, winreg.KEY_READ) as key:
        old_path = winreg.QueryValueEx(key, "Path")[0]
    cuda_bin = os.path.join(CUDA_PATH, "bin")
    if cuda_bin.lower() not in old_path.lower():
        new_path = f"{cuda_bin};{old_path}"
        subprocess.call(f'setx PATH "{new_path}" /M', shell=True)
        print(f"[信息] 已将 {cuda_bin} 添加到系统PATH")
    else:
        print(f"[信息] {cuda_bin} 已在系统PATH中")

def main():
    if not is_admin():
        print("[错误] 请以管理员身份运行本脚本！")
        sys.exit(1)
    install_cuda()
    install_cudnn()
    set_env()
    print("\n===== CUDA Toolkit 12.9.1 + cuDNN 安装与配置完成 =====")
    print("如需使用请重启命令行或注销/重启系统以刷新环境变量")
    input("按回车键退出...")

if __name__ == '__main__':
    main() 

# # 复制必须的CUDA DLL（根据你的实际CUDA版本调整）
# Copy-Item "$env:CUDA_PATH\bin\cudart64_12.dll" -Destination "D:\ComfyUINew\ComfyUI\python\"
# Copy-Item "$env:CUDA_PATH\bin\cublas64_12.dll" -Destination "D:\ComfyUINew\ComfyUI\python\"
# Copy-Item "$env:CUDA_PATH\bin\cudnn_ops_infer64_8.dll" -Destination "D:\ComfyUINew\ComfyUI\python\"  # 如果使用cuDNN